import os
import numpy as np
import pandas
from tqdm import tqdm

dataset_root = './../data'
dataset_name = 'DepthTrack'
gt_file_name = 'groundtruth.txt'
gt_split_char = ','

data_dic = {}

# circle all video
txt_file = os.path.join(dataset_root, dataset_name, 'depthtrack_val.txt')
with open(txt_file, 'r') as f:
    seq_names = f.read().splitlines()
    seq_names = [s.strip().split('.')[0] for s in seq_names]

for seq_name in tqdm(seq_names):
    seq_dir = os.path.join(dataset_root, dataset_name, seq_name)
    # gt file
    gt_file_path = os.path.join(seq_dir, gt_file_name)
    with open(gt_file_path, 'r') as f:
        gt_lines = f.read()
    gt_lines = gt_lines.replace('nan', '0', -1)
    gt_lines = gt_lines.splitlines()
    gt_lines = [s.strip().split(gt_split_char) for s in gt_lines]
    gt_lines = np.array(gt_lines, dtype=np.int32).tolist()
    # img path
    img_dir = os.path.join(seq_dir, 'color')
    img_list = [f for f in os.listdir(img_dir) if f.endswith('.png') or f.endswith('.jpg')]
    # sort
    img_list.sort(key=lambda x: int(x.split('.')[0]))
    img_list = [f'{seq_name}/color/{f}' for f in img_list]
    # attr
    attr = [attr.split('.')[0] for attr in os.listdir(seq_dir) if attr.endswith('.tag')]
    # absent
    absent = [1] * len(img_list)

    seq_dic = {
        'video_dir': seq_name,
        'init_rect': gt_lines[0],
        'img_names': img_list,
        'gt_rect': gt_lines,
        'attr': attr,
        'absent': absent
    }
    data_dic[seq_name] = seq_dic

# save to json
json_path = os.path.join(dataset_root, dataset_name, 'depthtrack_bike.json')
pandas.DataFrame(data_dic).to_json(json_path)